#include "faiss_ivfpq_index.h"
#include <faiss/index_io.h>
#include <filesystem>

FaissIVFPQIndex::FaissIVFPQIndex(const ConfigParser& config) {
    M = config.get_int("M", 8);
    k_f = config.get_int("k_f", 68);
    nbits = config.get_int("nbits", 8);
    nprobe = config.get_int("nprobe", 145);
#ifdef QUANT
    qbs = config.get_int("qbs", 32);
#endif
    std::string metric_str = config.get_string("metric", "L2");
    if (metric_str == "IP") {
        metric = faiss::METRIC_INNER_PRODUCT;
    } else {
        metric = faiss::METRIC_L2;
    }
    index_path = config.get_string("index_path", "index.faiss");
    save_or_load = config.get_string("save_or_load", "neither");
}

FaissIVFPQIndex::~FaissIVFPQIndex() {
    if (refine_index) {
        refine_index.reset();
    } else {
        index.reset();
    }
}

void FaissIVFPQIndex::build(const Datasets* data) {
    if (save_or_load == "load") {
        auto loaded_index = faiss::read_index(index_path.c_str());
        auto refine_ptr = dynamic_cast<faiss::IndexRefineFlat*>(loaded_index);
        if (refine_ptr) {
            refine_index.reset(refine_ptr);
            index.reset();
            quantizer.reset();
        } else {
            auto ivfpq_ptr = dynamic_cast<faiss::IndexIVFPQ*>(loaded_index);
            if (ivfpq_ptr) {
                index.reset(ivfpq_ptr);
                refine_index.reset();
                quantizer.reset();
                if (index) index->own_fields = true;
            } else {
                delete loaded_index;
                throw std::runtime_error("Loaded index is not IVFPQ or RefineFlat(IVFPQ)");
            }
        }
    } else {
        int nlist = int(4 * sqrt(data->nb));
        if (metric == faiss::METRIC_L2) {
            quantizer = std::make_unique<faiss::IndexFlatL2>(data->d);
        } else {
            quantizer = std::make_unique<faiss::IndexFlatIP>(data->d);
        }
        index = std::make_unique<faiss::IndexIVFPQ>(quantizer.release(), data->d, nlist, M, nbits, metric);
#ifdef QUANT
        index->quant_bits = qbs;
#endif
        refine_index = std::make_unique<faiss::IndexRefineFlat>(index.release());
        refine_index->train(data->nb, data->xb);
        refine_index->add(data->nb, data->xb);
        if (save_or_load == "save") {
			std::filesystem::path dir_path = std::filesystem::path(index_path).parent_path();
			if (!dir_path.empty() && !std::filesystem::exists(dir_path)) {
                std::filesystem::create_directories(dir_path);
            }
            faiss::write_index(refine_index.get(), index_path.c_str());
        }
    }
}
void FaissIVFPQIndex::prepare_search() {
    ivf_params.nprobe = nprobe;
    refine_params.base_index_params = &ivf_params;
    refine_params.k_factor = k_f;
}

void FaissIVFPQIndex::search_single(const float* xq, int k, float* distances, int64_t* labels) {
    int n = 1;
    refine_index->search(n, xq, k, distances, labels, &refine_params);
}

void FaissIVFPQIndex::search_batch(const int nq, const float* xq, int k, float* distances, int64_t* labels) {
    refine_index->search(nq, xq, k, distances, labels, &refine_params);
}